﻿# coding: UTF-8
# Name: 打开URL并获取源码
# Author: LYC
# Created:

import sys
import requests
import argparse

def encode(x):
    """对内容进行编码"""
    return x.encode(sys.getdefaultencoding(), u"ignore")

class Console(object):
    """
    控制台
    """
    enable = True
    output_prefix = u""
    intput_prefix = u""
    new_line = u"\n"

    @classmethod
    def write(cls, msg=u"", *args):
        """输出内容"""
        if args:msg = msg % args
        sys.stdout.write(encode(msg))

    @classmethod
    def write_line(cls, msg=u"", *args):
        """输出一行"""
        cls.write(u"%s%s%s" % (cls.output_prefix, msg, cls.new_line), *args)

    @classmethod
    def write_dict(cls, dct):
        """输出字典"""
        for k, v in dct.items():
            cls.write_line("\t%s = %s", k, v)

    @classmethod
    def read_line(cls, msg=u"", typ=None):
        """读取内容"""
        data = raw_input(u"%s%s" % (cls.intput_prefix, msg))
        if typ:data = typ(data)
        return data

def args_to_dict(args):
    """转为字典"""
    if isinstance(args,str):args = [args]
    return dict(i.split(u"=") for i in args if i)

def filter_html(query, html):
    """过滤内容"""
    import pyquery
    pq = pyquery.PyQuery(html)
    data = pq(query)

    lmd = lambda x:pyquery.PyQuery(x).outer_html()
    return encode(Console.new_line.join(map(lmd, data)))

def save_data(content, name):
    """保存数据"""
    import os
    if os.path.exists(name):
        d, n = os.path.split(name)

        name = "(%d)" + n
        for i in xrange(1, 65535):
            if not os.path.exists(name % i):
                name = os.path.join(d, name % i)
                break

    with open(name, u"wb") as fp:
        fp.write(content)
    return name

class UrlOpener(object):
    """
    Url访问器
    """

    headers = {}

    def __init__(self, proxy_lst=[], timeout=1):
        super(UrlOpener, self).__init__()
        self.headers = UrlOpener.headers
        self.proxy = args_to_dict(proxy_lst)
        self.timeout = timeout

    @classmethod
    def format_url(cls, url):
        """格式化url"""
        import urlparse
        url_split = list(urlparse.urlparse(url, "http"))
        if not url_split[1] and url_split[2]:
            url_split[1], url_split[2] = url_split[2], url_split[1]

        return urlparse.urlunparse(url_split)

    def add_headers(self, items):
        """添加请求头"""
        items = args_to_dict(items)
        self.headers.update(items)

    def get(self, url, data={}):
        """使用get方法"""
        url = self.format_url(url)
        return requests.get(url, params=data, headers=self.headers, proxies=self.proxy, timeout=self.timeout)

    def post(self, url, data={}):
        """使用post方法"""
        import json
        url = self.format_url(url)
        return requests.post(url, data=json.dumps(data), headers=self.headers, proxies=self.proxy, timeout=self.timeout)

def main(argv):
    parser = argparse.ArgumentParser(description=u"Get data from a url.")
    parser.add_argument(u"-u", u"--url", action=u"store", nargs=u"?", help=u"the url you want to open")
    parser.add_argument(u"-o", u"--out", action=u"store", default=u"", help=u"the file to output,stdout is default")
    parser.add_argument(u"-q", u"--query", action=u"store", default=u"", help=u"the jQuery type string to filter the html")
    parser.add_argument(u"-d", u"--data", action=u"store", nargs=u"*", default=[], help=u"the data you want to send")
    parser.add_argument(u"-m", u"--method", action=u"store", default=u"get", help=u"choice the method(get/post),get is default")
    parser.add_argument(u"-p", u"--proxy", action=u"store", default=[], help=u"open url by proxies")
    parser.add_argument(u"-t", u"--timeout", action=u"store", default=300, type=float, help=u"timeout")
    parser.add_argument(u"-hd", u"--header", action=u"store", nargs=u"*", default=[], help=u"add you request headers")

    args = parser.parse_args(argv[1:])

    if not args.url:
        args.url = Console.read_line(u"URL:")

    Console.write_line(u"URL: %s", args.url)
    Console.write_line(u"Timeout: %s seconds", args.timeout)

    opener = UrlOpener(args.proxy, args.timeout)
    if opener.proxy:
        Console.write_line(u"Proxies:")
        Console.write_dict(opener.proxy)

    opener.add_headers(args.header)
    if opener.headers:
        Console.write_line(u"Headers:")
        Console.write_dict(opener.headers)

    data = args_to_dict(args.data)
    if data:
        Console.write_line(u"Data:")
        Console.write_dict(data)

    method = args.method.lower()
    opener_method = opener.get
    if method == u"get":
        opener_method = opener.get
    elif method == u"post":
        opener_method = opener.post
    else:
        raise ValueError(u"cound not recognize the method of %s" % method)

    Console.write_line(u"Method: %s", method)
    Console.write_line()

    Console.write_line(u"Sending request...")

    response = opener_method(args.url, data)

    Console.write_line(u"Get a response from %s", response.url)

    Console.write_line(u"Status code: %s" % response.status_code)
    Console.write_line(u"Reason: %s" % response.reason)

    Console.write_line(u"Data length: %s", len(response.text))
    Console.write_line(u"Data encoding: %s", response.encoding)

    html = response.text
    content = response.content
    if args.query:
        Console.write_line(u"Query selector: %s", args.query)
        html = filter_html(args.query, html)
        content = html

    if args.out:
        name = save_data(content, args.out)

        Console.write_line(u"File name: %s", name)

    elif response.encoding:
        Console.write_line(u"Text:")
        Console.write_line(html)

    response.close()

if __name__ == '__main__':
    main(sys.argv)